/***********************************************************
     Starter code for Assignment 3

     This code was originally written by Jack Wang for
		    CSC418, SPRING 2005

		Implementations of functions in raytracer.h, 
		and the main function which specifies the 
		scene to be rendered.	

***********************************************************/


#include "raytracer.h"
#include "MotionBlur.h"
#include "Utils/bmp_io.h"
#include <cmath>
#include <iostream>
#include <cstdlib>
#include <string>
#include <map>
#include "Materials/SolidMaterial.h"
#include "Materials/TextureMaterial.h"
#include "Materials/UnitSquareTextureMapper.h"
#include "Shapes/UnitCylinderShape.h"
#include "Lights/FiniteLight.h"
#include "Lights/ClusterLight.h"
#include "Shapes/UnionShape.h"
#include "Shapes/UnitCubeShape.h"
#include "Shapes/IntersectionShape.h"
#include "Shapes/TriangleShape.h"
#include "Shapes/ModelFactory.h"
#include "Lights/FiniteLight.h"

#include "Scenes/TextureScene.h"
#include "Scenes/DefaultScene.h"
#include "Scenes/RefractionScene.h"
#include "Scenes/MotionBlurScene.h"
#include "Scenes/GlossyScene.h"
#include "Scenes/SoftShadowsScene.h"
#include "Scenes/FinalScene.h"
#include "Scenes/ClusterLightScene.h"
#include "Scenes/UnionShapeScene.h"
#include "Scenes/IntersectionShapeScene.h"
#include "Scenes/TriangleScene.h"
#include "Scenes/KettleScene.h"
#include "Scenes/CarScene.h"
#include "Scenes/FaceScene.h"
#include "Scenes/TextureTestScene.h"
#include "Scenes/BoxTextureTestScene.h"

// Rendering settings
bool signature = false;
bool shadow = true;
bool antiAlias = false;
int antiAliasSampleSize = 4;
bool enableSpecular = true;
double focalLen = -1.0;
std::string scene = "default";
unsigned int motionBlurIters = 1;
const unsigned int MotionBlurItersDefault = 10;
bool dofEnabled = false;
double depth = 2;
double aperture = 0.05;
unsigned int dofSampleSize = 1;
unsigned int DofSampleSizeDefault = 10;

Raytracer::Raytracer() : _lightSource(NULL), MAX_LEVEL(3), _pMotionBlur(NULL)
{
	_root = new SceneDagNode();
}

Raytracer::~Raytracer() {
	delete _root;
}

SceneDagNode* Raytracer::addObject(SceneDagNode* parent,
                                   SceneObject* obj) {
    
	SceneDagNode* node = new SceneDagNode(obj);
	node->parent = parent;
	node->next = NULL;
	node->child = NULL;
	
	// Add the object to the parent's child list, this means
	// whatever transformation applied to the parent will also
	// be applied to the child.
	if (parent->child == NULL) {
		parent->child = node;
	}
	else {
		parent = parent->child;
		while (parent->next != NULL) {
			parent = parent->next;
		}
		parent->next = node;
	}
	
	return node;;
}

LightListNode* Raytracer::addLightSource( LightSource* light, bool enableSpecular) {
	LightListNode* tmp = _lightSource;
	_lightSource = new LightListNode( light, tmp );
    light->enableSpecularShading(enableSpecular);
	return _lightSource;
}

Matrix4x4 Raytracer::initInvViewMatrix( Point3D eye, Vector3D view, 
		Vector3D up ) {
	Matrix4x4 mat; 
	Vector3D w;
	view.normalize();
	up = up - up.dot(view)*view;
	up.normalize();
	w = view.cross(up);

	mat[0][0] = w[0];
	mat[1][0] = w[1];
	mat[2][0] = w[2];
	mat[0][1] = up[0];
	mat[1][1] = up[1];
	mat[2][1] = up[2];
	mat[0][2] = -view[0];
	mat[1][2] = -view[1];
	mat[2][2] = -view[2];
	mat[0][3] = eye[0];
	mat[1][3] = eye[1];
	mat[2][3] = eye[2];

	return mat; 
}

void Raytracer::traverseScene( SceneDagNode* node, Ray3D& ray ) {
	SceneDagNode *childPtr;

	// Applies transformation of the current node to the global
	// transformation matrices.
	_modelToWorld = _modelToWorld*node->obj->getTransform();
	_worldToModel = node->obj->getInvTransform()*_worldToModel;
	if (node->obj) {
		// Perform intersection.
		node->obj->intersect(ray, _worldToModel, _modelToWorld);
	}
	// Traverse the children.
	childPtr = node->child;
	while (childPtr != NULL) {
		traverseScene(childPtr, ray);
		childPtr = childPtr->next;
	}

	// Removes transformation of the current node from the global
	// transformation matrices.
	_worldToModel = node->obj->getTransform()*_worldToModel;
	_modelToWorld = _modelToWorld*node->obj->getInvTransform();
}

void Raytracer::computeShading( Ray3D& ray ) {
	LightListNode* curLight = _lightSource;
	for (;;) {
		if (curLight == NULL) break;
		// Each lightSource provides its own shading function.

		// Implement shadows here if needed.

		curLight->light->shade(ray);
		curLight = curLight->next;
	}
}

void Raytracer::initPixelBuffer() {
	int numbytes = _scrWidth * _scrHeight * sizeof(unsigned char);
	_rbuffer = new unsigned char[numbytes];
	_gbuffer = new unsigned char[numbytes];
	_bbuffer = new unsigned char[numbytes];
	for (int i = 0; i < _scrHeight; i++) {
		for (int j = 0; j < _scrWidth; j++) {
			_rbuffer[i*_scrWidth+j] = 0;
			_gbuffer[i*_scrWidth+j] = 0;
			_bbuffer[i*_scrWidth+j] = 0;
		}
	}
}

void Raytracer::flushPixelBuffer( char *file_name ) {
	bmp_write( file_name, _scrWidth, _scrHeight, _rbuffer, _gbuffer, _bbuffer );
	delete _rbuffer;
	delete _gbuffer;
	delete _bbuffer;
}

void Raytracer::navigateScene(Ray3D& ray)
{
    traverseScene(_root, ray);
}

Colour Raytracer::shadeRay( Ray3D& ray )
{
    return shadeRay(ray, 1);
}	

// Returns global ambient color
Colour Raytracer::getAmbientColor(const Intersection& intersection, bool inShadow)
{
    // Compute shadows and if the pixel is not in direct range of any lights cast shadow
    LightListNode* current = _lightSource;
    
    Colour globalAmbientColour;
    
    const Material* mat = intersection.obj->getMaterial();
    
    while (current)
    {
        globalAmbientColour = globalAmbientColour + (current->light->getAmbientColour() * mat->getColour(Ambient, intersection, inShadow));
        current = current->next;
    }
    
    globalAmbientColour.clamp();
    return globalAmbientColour;
}

Colour Raytracer::shadeShadow(Ray3D& ray, bool isInShadow)
{
    // Compute shadows and if the pixel is not in direct range of any lights cast shadow
    LightListNode* current = _lightSource;
    
    Colour shadowColour;
    while (current)
    {
        shadowColour = shadowColour + current->light->shadeShadow(ray, *this, isInShadow);
        current = current->next;
    }
    
    return shadowColour;
}

// Determines whether a specified point is in shadow
bool Raytracer::isInShadow(Point3D point)
{
    // Compute shadows and if the pixel is not in direct range of any lights cast shadow
    LightListNode* current = _lightSource;
    
    while (current)
    {
        if (!current->light->isInShadow(point, *this))
        {
            return false;
        }
        
        current = current->next;
    }
    
    return true;
}

// Returns the gloss colour
Colour Raytracer::getGlossColour(Ray3D& incidentRay, unsigned int glossFactor, int level)
{
    // Compute reflection ray
    Ray3D reflectionRay;
    const double eps = 0.01;
    reflectionRay.dir = incidentRay.dir - 2 * (incidentRay.intersection.normal.dot(incidentRay.dir)) * incidentRay.intersection.normal;
    reflectionRay.origin = incidentRay.intersection.point + eps * reflectionRay.dir;
    reflectionRay.dir.normalize();
    
    // creating an orthonormal basis with origin = intersection of incident ray
    // and w = direction of primary reflection
    Vector3D up(reflectionRay.origin[0], reflectionRay.origin[1] + 1.0, reflectionRay.origin[2]);
    Matrix4x4 view2world = initInvViewMatrix(reflectionRay.origin, reflectionRay.dir, up);
    Matrix4x4 world2view = view2world.transpose();
    
    double r = 0;
    double g = 0;
    double b = 0;
    const double glossyAperture = 0.03;
    
    for(int i = 0; i < glossFactor; i++)
    {
		Vector3D reflectionRayInView = world2view * reflectionRay.dir;
            
        double epsX = (double)rand() / RAND_MAX - 0.5;
        reflectionRayInView[0] += epsX * glossyAperture;
			
        double epsY = (double)rand() / RAND_MAX - 0.5;
        reflectionRayInView[1] += epsY * glossyAperture;
			
        reflectionRay.dir = view2world * reflectionRayInView;
			
        if(incidentRay.intersection.normal.dot(reflectionRay.dir) >= 0)
        {
            reflectionRayInView[0] *= -1.0;
            reflectionRayInView[1] *= -1.0;
				
            reflectionRay.dir = view2world * reflectionRayInView;
        }
        
        reflectionRay.dir.normalize();
        traverseScene(_root, reflectionRay);
        
        if ((!reflectionRay.intersection.none) &&
            (reflectionRay.intersection.obj->getMaterial()->getSpecularReflectivityFactor() > 0) &&
            incidentRay.intersection.normal.dot(reflectionRay.dir) >= 0)
        {
            Colour reflectionColour = shadeRay(reflectionRay, level+1);
            reflectionColour = incidentRay.intersection.obj->getMaterial()->getSpecularReflectivityFactor() * reflectionColour;
            r += reflectionColour[0];
            g += reflectionColour[1];
            b += reflectionColour[2];
        }
    }
    
    return Colour(r / (double)glossFactor, g / (double)glossFactor, b / (double)glossFactor);
}


// Returns the reflection colour
Colour Raytracer::getReflectionColor(Ray3D& incidentRay, int level)
{
    // Compute reflection ray
    Ray3D reflectionRay;
    const double eps = 0.01;
    reflectionRay.dir = incidentRay.dir - 2 * (incidentRay.intersection.normal.dot(incidentRay.dir)) * incidentRay.intersection.normal;
    reflectionRay.origin = incidentRay.intersection.point + eps * reflectionRay.dir;
    reflectionRay.dir.normalize();
    
    unsigned int glossFactor = incidentRay.intersection.obj->getMaterial()->getGlossFactor();
    
    if (glossFactor)
    {
        return getGlossColour(incidentRay, glossFactor, level);
    }
    
    traverseScene(_root, reflectionRay);
    
    if ((!reflectionRay.intersection.none) &&
        (reflectionRay.intersection.obj->getMaterial()->getSpecularReflectivityFactor() > 0))
    {
        Colour reflectionColor = shadeRay(reflectionRay, level+1);
        return incidentRay.intersection.obj->getMaterial()->getSpecularReflectivityFactor() * reflectionColor;
    }
    
    return Colour();
}

// Returns the refraction colour
Colour Raytracer::getRefractionColour(Ray3D& incidentRay, int level)
{
    // If material has no refractive index
    // then it does not support refraction
    const double n1 = 1;
    const double n2 = incidentRay.intersection.obj->getMaterial()->getRefractiveIndex();
    double n = 0;
    
    Colour refractionColour;
    if ((n2 == 0) || (incidentRay.intersection.obj->getMaterial()->getDiffuseReflectivityFactor() <= 0))
    {
        return refractionColour;
    }
    
    if (incidentRay.intersection.normal.dot(incidentRay.dir) < 0)
    {
        n = n1 / n2;
    }
    else
    {
        n = n2 / n1;
    }
    
    Vector3D normal = incidentRay.intersection.normal;
    normal.normalize();
    
    Vector3D incident = incidentRay.dir;
    incident.normalize();
    
    const double cosI = -1 * normal.dot(incident);
    const double sinT2 = n * n * (1.0 - (cosI * cosI));
    const double cosT2 = 1.0 - sinT2;
    
    if (sinT2 > 1)
    {
        return refractionColour;
    }
    
    // Compute refraction ray
    const double eps = 0.1;
    Ray3D refractionRay;
    refractionRay.dir = n * incident - (n * cosI - sqrtf(cosT2)) * normal;
    refractionRay.dir.normalize();
    refractionRay.origin = incidentRay.intersection.point + eps * refractionRay.dir;
    
    traverseScene(_root, refractionRay);
    
    if (!refractionRay.intersection.none)
    {
        refractionColour = shadeRay(refractionRay, level+1);
        Colour absorbance = refractionRay.intersection.t_value *
                            incidentRay.intersection.obj->getMaterial()->getDiffuseReflectivityFactor() *
                            incidentRay.intersection.obj->getMaterial()->getColour(Ambient, incidentRay.intersection, false);
        Colour transparency = Colour(exp(-1 * absorbance[0]), exp(-1 * absorbance[1]), exp(-1 * absorbance[2]));        
        refractionColour = transparency * refractionColour;
    }        
    
    return refractionColour;
}

Colour Raytracer::shadeRay( Ray3D& ray, int level)
{
    Colour col(0.0, 0.0, 0.0);
	
    traverseScene(_root, ray);
    
	// Don't bother shading if the ray didn't hit
	// anything.
	if (ray.intersection.none)
    {
        return col;
    }    
    
    if (signature)
    {
        double colorValue = 1.0 / (double)ray.intersection.obj->getShapeId();
        return Colour(colorValue, colorValue, colorValue);
    }

    if (level == MAX_LEVEL)
    {
        col.clamp();
        return col;
    }
    
    bool inShadow = isInShadow(ray.intersection.point);
    col = getAmbientColor(ray.intersection, inShadow);
    
    // Compute color of ray before shadow is applied
    computeShading(ray);
    
    if (shadow)
    {
        col = col + shadeShadow(ray, inShadow);
        
        if (inShadow)
        {
            // If the point is in shadow, do not recurse
            return col;
        }
    }
    else
    {
        // Shadows are disabled
        col = col + ray.col;
    }
    
    // Compute reflection ray
    col = col + getReflectionColor(ray, level);
    
    // Compute refraction ray
    col = col + getRefractionColour(ray, level);
    
    col.clamp();

	return col;
}

void Raytracer::render( int width, int height, Point3D eye, Vector3D view,
		Vector3D up, double fov, char* fileName ) {
	Matrix4x4 viewToWorld;
	_scrWidth = width;
	_scrHeight = height;
	double factor = (double(height)/2)/tan(fov*M_PI/360.0);

	initPixelBuffer();
	viewToWorld = initInvViewMatrix(eye, view, up);
        
    // Construct a ray for each pixel.
    for (int i = 0; i < _scrHeight; i++)
    {
        for (int j = 0; j < _scrWidth; j++)
        {
            Point3D imagePlane;

            double r = 0;
            double g = 0;
            double b = 0;
            int n = antiAlias ? antiAliasSampleSize : 1;
                
            for (int p = 0; p < n; p++)
            {
                for (int q = 0; q < n; q++)
                {
                    for (int dofIter = 0; dofIter < dofSampleSize; dofIter++)
                    {
                        for (int motionBlurIter = 0; motionBlurIter < motionBlurIters; motionBlurIter++)
                        {
                            if (_pMotionBlur)
                            {
                                _pMotionBlur->startMotionBlur();
                            }
                    
                            double epsX = (double)rand() / (double)RAND_MAX;
                            double epsY = (double)rand() / (double)RAND_MAX;
                            imagePlane[0] = (-double(width)/2 + (antiAlias ? (p + epsX) / 4 : 0.5) + j)/factor;
                            imagePlane[1] = (-double(height)/2 + (antiAlias ? (q + epsY) / 4  : 0.5)+ i)/factor;
                            imagePlane[2] = -1;

                            // Sets up ray origin and direction in view space,
                            // image plane is at z = -1.
                            Point3D origin(0, 0, 0);
                            
                            Ray3D ray;
                            ray.dir = (imagePlane - origin);
                            ray.dir.normalize();                            
                        
                            // If dof is enabled then we find the point on the focal plane, fix that and then take
                            // random points around the origin within the aperture and treat that as the origin
                            if (dofEnabled)
                            {                                
                                Point3D pointOnFocalPlane = origin + (-1 * depth / ray.dir[2]) * ray.dir;
                                
                                // Now that we know the point of the focal plane, fix it and pick a point within the
                                // aperteur around the origin and take it as the new origin
                                double xRand = (double)rand() / (double)RAND_MAX;
                                origin[0] = origin[0] + (0.5 - xRand) * aperture;
                                
                                double yRand = (double)rand() / (double)RAND_MAX;
                                origin[1] = origin[1] + (0.5 - yRand) * aperture;
                                
                                ray.dir = pointOnFocalPlane - origin;
                                ray.dir.normalize();
                            }
                            
                            ray.origin = viewToWorld * origin;
                            ray.dir = viewToWorld * ray.dir;
                            
                            Colour shadeColour = shadeRay(ray);
                        
                            r += shadeColour[0];
                            g += shadeColour[1];
                            b += shadeColour[2];
                        
                            if (_pMotionBlur)
                            {
                                _pMotionBlur->endMotionBlur();
                            }
                        }
                    }
                }
            }
            
            double avgFactor = ((double)(1.0/(n * n * motionBlurIters * dofSampleSize)));
            Colour col(r * avgFactor, g * avgFactor, b * avgFactor);
            col.clamp();
                
            _rbuffer[i*width+j] = int(col[0]*255);
            _gbuffer[i*width+j] = int(col[1]*255);
            _bbuffer[i*width+j] = int(col[2]*255);
        }
    }

	flushPixelBuffer(fileName);
}

void Raytracer::setMotionBlur(MotionBlur *motionBlur)
{
    _pMotionBlur = motionBlur;
}

int main(int argc, char* argv[])
{
	// Build your scene and setup your camera here, by calling
	// functions from Raytracer.  The code here sets up an example
	// scene and renders it from two different view points, DO NOT
	// change this if you're just implementing part one of the 
	// assignment.  
	Raytracer raytracer;
	int width = 320; 
	int height = 240;
    
	if (argc > 2) {
		width = atoi(argv[1]);
		height = atoi(argv[2]);
	}
    
    for (int i = 3; i < argc; i++)
    {
        std::string arg = std::string(argv[i]);
        
        // Turn on signature mode
        if (arg == "-signature")
        {
            signature = true;
        }
        else if (arg == "-disable_specular")
        {
            enableSpecular = false;
        }
        else if (arg == "-motion_blur")
        {
            double motionBlurArg = 0;
            if ((i < argc - 1) && ((motionBlurArg = atoi(argv[i+1])) > 0))
            {
                motionBlurIters = motionBlurArg;
                i++;
            }
            else
            {
                motionBlurIters = MotionBlurItersDefault;
            }
        }
        else if (arg == "-disable_shadows")
        {
            shadow = false;
        }
        else if (arg == "-anti_alias")
        {
            antiAlias = true;
            int argAntiAliasSampleSize = 0;
            
            if ((i < argc - 1) && ((argAntiAliasSampleSize = atoi(argv[i+1])) > 0))
            {
                antiAliasSampleSize = argAntiAliasSampleSize;
                i++;
            }
        }       
        else if (arg == "-scene")
        {
            if (i < argc - 1)
            {
                scene = argv[i+1];
                i++;
            }
        }
        else if (arg == "-dof")
        {
            dofEnabled = true;
            
            double argDepth = 0;
            
            if ((i < argc - 1) && ((argDepth = atof(argv[i+1])) > 0))
            {
                depth = argDepth;
                i++;
            }
            
            double argAperture = 0;
            
            if ((i < argc - 1) && ((argAperture = atof(argv[i+1])) > 0))
            {
                aperture = argAperture;
                i++;
            }
            
            double argDofSampleSize = 0;
            if ((i < argc - 1) && ((argDofSampleSize = atoi(argv[i+1])) > 0))
            {
                dofSampleSize = argDofSampleSize;
                i++;
            }
            else
            {
                dofSampleSize = DofSampleSizeDefault;
            }            
        }
    }
    
    typedef void (*SceneFunction)(Raytracer&, int, int, bool);
    std::map<std::string, SceneFunction> scenes;
    scenes["texture"] = &renderTextureScene;
    scenes["default"] = &renderDefaultScene;
    scenes["motionblur"] = &renderMotionBlurScene;
    scenes["refraction"] = &renderRefractionScene;
    scenes["glossy"] = &renderGlossyScene;
    scenes["softshadow"] = &renderSoftShadowScene;
    scenes["final"] = &renderFinalScene;
    scenes["cluster"] = &renderClusterLightScene;
    scenes["union"] = &renderUnionShapeScene;
    scenes["intersection"] = &renderIntersectionShapeScene;
    scenes["triangle"] = &renderTriangleScene;
    scenes["kettle"] = &renderKettleScene;
    scenes["car"] = &renderCarScene;
    scenes["face"] = &renderFaceScene;
    scenes["texture_test"] = &renderTextureTestScene;
    scenes["box_texture_test"] = &renderBoxTextureTestScene;
    
    std::map<std::string, SceneFunction>::iterator sceneIter = scenes.find(scene);
    
    if (sceneIter == scenes.end())
    {
        printf("The scene %s does not exist.", scene.c_str());
        return -1;
    }
    
    (*sceneIter->second)(raytracer, width, height, enableSpecular);
    
	return 0;
}

